package com.yisin.ssh2;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;

import com.yisin.ssh2.jsch.Channel;
import com.yisin.ssh2.jsch.ChannelExec;
import com.yisin.ssh2.jsch.JSchException;
import com.yisin.ssh2.jsch.Util;
import com.yisin.ssh2.stream.CombinedStream;

public class Scp extends SshTransferProtocolBase {

	private boolean m_recursive = false;
	private boolean m_verbos = false;
	private boolean m_cancelled = false;

	public Scp(String host, String user, String password) {
		super(host, user, password);
	}

	public Scp(String host, String user) {
		super(host, user);
	}

	protected String getChannelType() {
		return "exec";
	}

	// / <summary>
	// /This function is empty, so no channel is connected
	// /on session connect
	// / </summary>
	protected void ConnectChannel() {
	}

	// / <summary>
	// / Gets or sets a value indicating the default recursive transfer
	// behaviour
	// / </summary>
	public void setRecursive(boolean value) {
		m_recursive = value;
	}

	public boolean getRecursive() {
		return m_recursive;
	}

	// / <summary>
	// / Gets or sets a value indicating whether trace information should be
	// printed
	// / </summary>
	public void setVerbos(boolean value) {
		m_verbos = value;
	}

	public boolean getVerbos() {
		return m_verbos;
	}

	public void Cancel() {
		m_cancelled = true;
	}

	// / <summary>
	// / Creates a directory on the remot server
	// / </summary>
	// / <param name="dir">The new directory</param>
	@SuppressWarnings("unused")
	public void Mkdir(String dir) throws Exception {
		SCP_CheckConnectivity();

		Channel channel = null;
		CombinedStream server = null;
		m_cancelled = false;

		SCP_ConnectTo(channel, server, dir, true);
		SCP_EnterIntoDir(server, dir);
		if (null != channel) {
			channel.disconnect();
		}
	}

	public void Put(String fromFilePath, String toFilePath) throws Exception {
		this.To(fromFilePath, toFilePath);
	}

	public void Get(String fromFilePath, String toFilePath) throws Exception {
		this.From(fromFilePath, toFilePath);
	}

	// / <summary>
	// / Copies a file from local machine to a remote SSH machine.
	// / </summary>
	// / <param name="localPath">The local file path.</param>
	// / <param name="remotePath">The path of the remote file.</param>
	public void To(String localPath, String remotePath) throws Exception {
		this.To(localPath, remotePath, getRecursive());
	}

	// / <summary>
	// / Copies a file from local machine to a remote SSH machine.
	// / </summary>
	// / <param name="localPath">The local file path.</param>
	// / <param name="remotePath">The path of the remote file.</param>
	public void To(String localPath, String remotePath, boolean _recursive) throws Exception {
		SCP_CheckConnectivity();

		Channel channel = null;
		CombinedStream server = null;
		m_cancelled = false;

		try {
			// if we are sending a single file
			if (new File(localPath).exists()) {
				SCP_ConnectTo(channel, server, remotePath, _recursive);
				SCP_SendFile(server, localPath, remotePath);
				channel.disconnect();
			}
			// else, if we are sending a local directory
			else if (new File(localPath).exists()) {
				if (!_recursive) {
					throw new SshTransferException("'" + new File(localPath).getName() + "' is a directory, you should use recursive transfer.");
				}
				SCP_ConnectTo(channel, server, remotePath, true);
				ToRecursive(server, localPath, remotePath);
				channel.disconnect();
			} else {
				throw new SshTransferException("File not found: " + localPath);
			}
		} catch (Exception e) {
			if (getVerbos())
				System.err.println("Error: " + e.getMessage());
			// SendEndMessage(remoteFile, localFile, filesize,filesize,
			// "Transfer ended with an error.");
			try {
				channel.disconnect();
			} catch (Exception ex) {
			}
			throw e;
		}
	}

	// / <summary>
	// / Copies files and directories from local machine to a remote SSH machine
	// using SCP.
	// / </summary>
	// / <param name="server">I/O Stream for the remote server</param>
	// / <param name="src">Source to copy</param>
	// / <param name="dst">Destination path</param>
	private void ToRecursive(CombinedStream server, String src, String dst) throws IOException, SshTransferException {
		File srcFile = new File(src);
		File dstFile = new File(dst);
		if (srcFile.exists() && srcFile.isDirectory()) {
			SCP_EnterIntoDir(server, dstFile.getName());
			for (File file : srcFile.listFiles()) {
				if (file.isFile()) {
					SCP_SendFile(server, file.getAbsolutePath(), file.getName());
				}
			}
			if (m_cancelled) {
				return;
			}

			for (File file : srcFile.listFiles()) {
				if (file.isDirectory()) {
					ToRecursive(server, file.getAbsolutePath(), file.getName());
				}
			}
			SCP_EnterIntoParent(server);
		} else if (srcFile.exists() && srcFile.isFile()) {
			SCP_SendFile(server, src, srcFile.getName());
		} else {
			throw new SshTransferException("File not found: " + src);
		}
	}

	// / <summary>
	// / Copies a file from a remote SSH machine to the local machine using SCP.
	// / </summary>
	// / <param name="remoteFile">The remmote file name</param>
	// / <param name="localPath">The local destination path</param>
	public void From(String remoteFile, String localPath) throws Exception {
		this.From(remoteFile, localPath, getRecursive());
	}

	// / <summary>
	// / Copies a file from a remote SSH machine to the local machine using SCP.
	// / </summary>
	// / <param name="remoteFile">The remmote file name</param>
	// / <param name="localPath">The local destination path</param>
	// / <param name="recursive">Value indicating whether a recursive transfer
	// should take place</param>
	public void From(String remoteFile, String localPath, boolean _recursive) throws Exception {
		SCP_CheckConnectivity();

		Channel channel = null;
		CombinedStream server = null;
		m_cancelled = false;
		int filesize = 0;
		String filename = null;
		String cmd = null;

		try {
			String dir = null;
			if (new File(localPath).exists()) {
				dir = new File(localPath).getAbsolutePath();
			}

			SCP_ConnectFrom(channel, server, remoteFile, _recursive);

			byte[] buf = new byte[1024];

			// send '\0'
			SCP_SendAck(server);
			int c = SCP_CheckAck(server);

			// parse scp commands
			while ((c == 'D') || (c == 'C') || (c == 'E')) {
				if (m_cancelled)
					break;

				cmd = "" + (char) c;
				if (c == 'E') {
					c = SCP_CheckAck(server);
					dir = new File(dir).getName();
					if (getVerbos())
						System.out.println("E");
					// send '\0'
					SCP_SendAck(server);
					c = (char) SCP_CheckAck(server);
					continue;
				}

				// read '0644 ' or '0755 '
				server.read(buf, 0, 5);
				for (int i = 0; i < 5; i++)
					cmd += (char) buf[i];

				// reading file size
				filesize = 0;
				while (true) {
					server.read(buf, 0, 1);
					if (buf[0] == ' ')
						break;
					filesize = filesize * 10 + (buf[0] - '0');
				}

				// reading file name
				for (int i = 0;; i++) {
					server.read(buf, i, 1);
					if (buf[i] == (byte) 0x0a) {
						filename = Util.getString(buf, 0, i);
						break;
					}
				}
				cmd += " " + filesize + " " + filename;
				// send '\0'
				SCP_SendAck(server);

				// Receive file
				if (c == 'C') {
					if (getVerbos())
						System.out.println("Sending file modes: " + cmd);
					SCP_ReceiveFile(server, remoteFile, dir == null ? localPath : dir + "/" + filename, filesize);

					if (m_cancelled)
						break;

					// send '\0'
					SCP_SendAck(server);
				}
				// Enter directory
				else if (c == 'D') {
					if (dir == null) {
						if (new File(localPath).exists())
							throw new SshTransferException("'" + localPath + "' is not a directory");
						dir = localPath;
						new File(dir).mkdir();
					}
					if (getVerbos())
						System.out.println("Entering directory: " + cmd);
					dir += "/" + filename;
					new File(dir).mkdir();
				}

				c = SCP_CheckAck(server);
			}
			channel.disconnect();
		} catch (Exception e) {
			if (getVerbos())
				e.printStackTrace();
			try {
				channel.disconnect();
			} catch (Exception ex) {
			}
			throw e;
		}
	}

	// / <summary>
	// / Checks is a channel is already connected by this instance
	// / </summary>
	protected void SCP_CheckConnectivity() throws Exception {
		if (!Connected())
			throw new Exception("Channel is down.");
	}

	// / <summary>
	// / Connect a channel to the remote server using the 'SCP TO' command ('scp
	// -t')
	// / </summary>
	// / <param name="channel">Will contain the new connected channel</param>
	// / <param name="server">Will contaun the new connected server I/O
	// stream</param>
	// / <param name="rfile">The remote path on the server</param>
	// / <param name="recursive">Idicate a recursive scp transfer</param>
	protected void SCP_ConnectTo(Channel channel, CombinedStream server, String rfile, boolean recursive) throws IOException, JSchException, SshTransferException {
		String scpCommand = "scp -p -t ";
		if (recursive)
			scpCommand += "-r ";
		scpCommand += "\"" + rfile + "\"";

		channel = (ChannelExec) m_session.openChannel(getChannelType());
		((ChannelExec) channel).setCommand(scpCommand);

		server = new CombinedStream(channel.getInputStream(), channel.getOutputStream());
		channel.connect();

		SCP_CheckAck(server);
	}

	// / <summary>
	// / Connect a channel to the remote server using the 'SCP From' command
	// ('scp -f')
	// / </summary>
	// / <param name="channel">Will contain the new connected channel</param>
	// / <param name="server">Will contaun the new connected server I/O
	// stream</param>
	// / <param name="rfile">The remote path on the server</param>
	// / <param name="recursive">Idicate a recursive scp transfer</param>
	protected void SCP_ConnectFrom(Channel channel, CombinedStream server, String rfile, boolean recursive) throws IOException, JSchException {
		String scpCommand = "scp -f ";
		if (recursive)
			scpCommand += "-r ";
		scpCommand += "\"" + rfile + "\"";

		channel = (ChannelExec) m_session.openChannel(getChannelType());
		((ChannelExec) channel).setCommand(scpCommand);

		server = new CombinedStream(channel.getInputStream(), channel.getOutputStream());
		channel.connect();

		// SCP_CheckAck(server);
	}

	// / <summary>
	// / Transfer a file to the remote server
	// / </summary>
	// / <param name="server">A connected server I/O stream</param>
	// / <param name="src">The source file to copy</param>
	// / <param name="dst">The remote destination path</param>
	protected void SCP_SendFile(CombinedStream server, String src, String dst) throws IOException, SshTransferException {
		int filesize = 0;
		int copied = 0;

		filesize = (int) (new File(src)).length();

		byte[] tmp = new byte[1];

		// send "C0644 filesize filename", where filename should not include '/'

		String command = "C0644 " + filesize + " " + new File(dst).getName() + "\n";
		if (getVerbos())
			System.out.println("Sending file modes: " + command);
		SendStartMessage(src, dst, filesize, "Starting transfer.");
		byte[] buff = Util.getBytes(command);
		server.write(buff, 0, buff.length);
		server.flush();

		if (SCP_CheckAck(server) != 0) {
			throw new SshTransferException("Error openning communication channel.");
		}

		// send a content of lfile
		SendProgressMessage(src, dst, copied, filesize, "Transferring...");
		FileInputStream fis = new FileInputStream(src);
		byte[] buf = new byte[1024 * 10 * 2];

		while (!m_cancelled) {
			int len = fis.read(buf, 0, buf.length);
			if (len <= 0)
				break;
			server.write(buf, 0, len);
			server.flush();
			copied += len;
			SendProgressMessage(src, dst, copied, filesize, "Transferring...");
		}
		fis.close();

		if (m_cancelled)
			return;

		// send '\0'
		buf[0] = 0;
		server.write(buf, 0, 1);
		server.flush();

		SendProgressMessage(src, dst, copied, filesize, "Verifying transfer...");
		if (SCP_CheckAck(server) != 0) {
			SendEndMessage(src, dst, copied, filesize, "Transfer ended with an error.");
			throw new SshTransferException("Unknow error during file transfer.");
		}
		SendEndMessage(src, dst, copied, filesize, "Transfer completed successfuly (" + copied + " bytes).");
	}

	// / <summary>
	// / Transfer a file from the remote server
	// / </summary>
	// / <param name="server">A connected server I/O stream</param>
	// / <param name="rfile">The remote file to copy</param>
	// / <param name="lfile">The local destination path</param>
	protected void SCP_ReceiveFile(CombinedStream server, String rfile, String lfile, int size) throws IOException, SshTransferException {
		int copied = 0;
		SendStartMessage(rfile, lfile, size, "Connected, starting transfer.");
		// read a content of lfile
		FileOutputStream fos = new FileOutputStream(lfile);
		int foo;
		int filesize = size;
		byte[] buf = new byte[1024];
		while (!m_cancelled) {
			if (buf.length < filesize)
				foo = buf.length;
			else
				foo = filesize;
			int len = server.read(buf, 0, foo);
			copied += len;
			fos.write(buf, 0, foo);
			SendProgressMessage(rfile, lfile, copied, size, "Transferring...");
			filesize -= foo;
			if (filesize == 0)
				break;
		}
		fos.close();
		if (m_cancelled)
			return;
		SCP_CheckAck(server);
		SendEndMessage(rfile, lfile, copied, size, "Transfer completed successfuly (" + filesize + " bytes).");
	}

	// / <summary>
	// / Instructs the remote server to enter into a directory
	// / </summary>
	// / <param name="server">A connected server I/O stream</param>
	// / <param name="dir">The directory name/param>
	protected void SCP_EnterIntoDir(CombinedStream server, String dir) {
		try {
			byte[] tmp = new byte[1];
			// send "C0644 filesize filename", where filename should not include
			// '/'

			String command = "D0755 0 " + new File(dir).getName() + "\n";
			if (getVerbos())
				System.out.println("Enter directory: " + command);

			byte[] buff = Util.getBytes(command);
			server.write(buff, 0, buff.length);
			server.flush();

			if (SCP_CheckAck(server) != 0) {
				throw new SshTransferException("Error openning communication channel.");
			}
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	// / <summary>
	// / Instructs the remote server to go up one level
	// / </summary>
	// / <param name="server">A connected server I/O stream</param>
	protected void SCP_EnterIntoParent(CombinedStream server) {
		try {
			byte[] tmp = new byte[1];

			// send "C0644 filesize filename", where filename should not include
			// '/'

			String command = "E\n";
			if (getVerbos())
				System.out.println(command);

			byte[] buff = Util.getBytes(command);
			server.write(buff, 0, buff.length);
			server.flush();

			if (SCP_CheckAck(server) != 0) {
				throw new SshTransferException("Error openning communication channel.");
			}
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	// / <summary>
	// / Gets server acknowledgment
	// / </summary>
	// / <param name="ins">A connected server I/O stream</param>
	private int SCP_CheckAck(CombinedStream ins) throws SshTransferException, IOException {
		int b = ins.readByte();
		// b may be 0 for success,
		// 1 for error,
		// 2 for fatal error,
		// -1
		if (b == 0)
			return b;
		if (b == -1)
			return b;

		if (b == 1 || b == 2) {
			StringBuilder sb = new StringBuilder();
			int c;
			do {
				c = ins.readByte();
				sb.append((char) c);
			} while (c != '\n');
			if (b == 1) { // error
				throw new SshTransferException(sb);
			}
			if (b == 2) { // fatal error
				throw new SshTransferException(sb);
			}
		}
		return b;
	}

	// / <summary>
	// / Sends acknowledgment to remote server
	// / </summary>
	// / <param name="server">A connected server I/O stream</param>
	private void SCP_SendAck(CombinedStream server) throws IOException {
		server.writeByte(0);
		server.flush();
	}

}
